Add tiled VAE parallelism & extended model support - #14

PatchGroupNorm summed each rank's variance about that rank's own mean, so the
result was short of the variance the unsharded norm computes by exactly how far
the patches sat from each other. A rank holding a brighter patch measured its
deviations from a brighter middle, and nothing put that offset back.

It also took torch.var at its default, which applies Bessel's correction, where
nn.GroupNorm divides by the count. That one is small enough to hide under a
loose tolerance in a single layer and accumulates through a decoder's worth.

Both are why a sharded AutoencoderKL decode did not reproduce a single-rank one.
The new test compares against nn.GroupNorm directly, over 4D and 5D inputs and
both split axes, at a tolerance that would have caught either.

Co-authored-by: Cursor <cursoragent@cursor.com>

DecoderAdapter is what every AutoencoderKL model shards through, xDiT's SD3 and
Z-Image among them, and its only check lived in a torchrun script that needs
NCCL and a GPU. WanDecoderAdapter had none at all. Both now decode against an
unsharded reference over gloo at world size 1, 2 and 4.

The harness these rest on gets its own test, since a comparison that quietly
accepted anything would turn the whole suite green and mean nothing.

Co-authored-by: Cursor <cursoragent@cursor.com>

Qwen-Image's decoder is a fork of Wan's: QwenImageCausalConv3d and
QwenImageResidualBlock are the Wan classes under other names, down to the same
causal padding tuple and the same RMS norms that need no sharding. The single
difference that reaches the adapter is that its up blocks were not given Wan's
first_chunk argument.

So rather than a second copy of five adapters, each Wan adapter now names the
block types it accepts and the child adapters it builds, and the Qwen ones
declare their own. The names xDiT dispatches on are unchanged, and the existing
Wan decode test pins the behaviour across the move.

Unlocks --use_parallel_vae for Qwen-Image, Qwen-Image-Edit, Krea-2-Raw and
Krea-2-Turbo. The blocks are imported optionally, so a diffusers too old to
carry them still runs every other family's adapter.

Co-authored-by: Cursor <cursoragent@cursor.com>

The direct path has two ways of padding a patch. For zeros it pads both sides
of the halo-extended input and crops the halo off the output, letting the
spurious padding at the interior edge land only on rows it then discards. For
any other mode it instead drops the padding at the interior edge, since the
halo already holds those rows, and the output is patch-sized as it stands. It
then cropped the halo off that too, taking a row off each interior boundary and
returning a patch shorter than the one it was given.

Deferring to build_crop_slice settles it: that already recognises an output the
size of the patch and leaves it alone, which is why the chunked path was right
all along.

Circular padding is now refused across ranks rather than answered wrongly. It
reads from the opposite edge of the image, which is not on a neighbour, so no
halo exchange can supply it. Nothing in scope uses it; a single rank still can.

Co-authored-by: Cursor <cursoragent@cursor.com>

…ecoders

Both pad their causal convolutions by replication, which left alone would have
each rank repeat its own edge rows at a boundary that is not an edge of the
image. Moving the spatial half of that padding into PatchConv3d fixes it, and
the temporal half stays where it was, one-sided along an axis nobody splits.

Where they differ is the two places sharding notices. HunyuanVideo normalises
with GroupNorm, whose statistics span the split and so have to be summed across
ranks, while 1.5 uses RMS and needs nothing. And HunyuanVideo's mid block
flattens (F, H, W) into a sequence and builds the causal mask itself, both from
the height in front of it, so its attention cannot be wrapped on its own the
way every other family's can: hand it a patch and it gets a mask cut for a
patch. Gathering around the whole mid block avoids reimplementing that forward
and costs little, since it runs at the latent resolution. 1.5's attention takes
the tensor whole and needs no such thing.

Unlocks --use_parallel_vae for HunyuanVideo, HunyuanVideo 1.5 and its distilled
variants.

Co-authored-by: Cursor <cursoragent@cursor.com>

LTX-2 is the last of the four families xDiT could not decode in parallel, and the least work
of them. Its causal convolution keeps spatial padding inside its nn.Conv3d rather than applying
it around one, so swapping that convolution for a PatchConv3d is the whole of the change; the
temporal padding repeats frames along an axis nobody splits. Its upsampler moves channels into
space, one input position per output, and its norms reduce over channels, so neither needs to
see rows a rank does not hold. Alone among these families its mid block has no attention, so
nothing has to be gathered back together mid-decode.

A residual block with inject_noise enabled is refused rather than sharded: the noise is drawn
per rank and would not add up to what a single rank draws, so the decode could not match its
reference. No shipped LTX-2 or LTX-2.3 config turns it on.

The decoder takes a timestep embedding where the other families take a temporal cache, so the
splitting and reassembling around the decode moves into _sharded_decode, which takes the call
to make rather than assuming a signature.

Co-authored-by: Cursor <cursoragent@cursor.com>

The QwenImage, HunyuanVideo and LTX-2 adapters name classes newer than this, which reads like
the floor is stale. It is not: those classes are resolved through diffusers_blocks rather than
imported, so only Wan's blocks have to exist for this package to import at all.

Co-authored-by: Cursor <cursoragent@cursor.com>

Patchify padded the axis it splits up to a multiple of the rank count, and the
adapters cropped the result afterwards. That is not the same computation as the
one being reproduced: the pad is zeros only until the first convolution, after
which it is the network's answer to zeros, and it reaches the rows that survive
the crop through every receptive field and every attention that follows.

Measured on a Qwen-Image decoder against a single-rank reference, 15 latent rows
over 2 ranks moved 84% of the output pixels by up to 0.77, and 16 rows over 3
moved 93% by up to 1.0. Wan's decoder, which has shipped for longer, is the same.
The 2D DecoderAdapter escaped it by splitting after its mid block rather than
before, so nothing that goes through AutoencoderKL is affected.

Bands are now cut in whole multiples of what the VAE narrows or widens the axis
by, which keeps each one on the grid the strided convolutions step along, and
uneven counts are absorbed by giving the first ranks one band more rather than by
inventing rows. Bands therefore differ in size between ranks, which the two
gathers could not do: dist.all_gather requires every rank to contribute the same
shape, so they now pad for the length of the transfer and slice on the far side,
where the padding cannot reach a convolution. That is also what WanZeroPadConv2d
was refusing uneven bands over, and it no longer needs to.

Each adapter's tests gain the case that was wrong. use_uniform_patch is gone from
the adapters, having named the behaviour that has been removed.

Co-authored-by: Cursor <cursoragent@cursor.com>

Qwen-Image's encoder is Wan 2.1's laid out flat and renamed, down to the resample
that pads (0, 1, 0, 1) and strides 2 over it. So the adapter is Wan's with a
different set of block classes, and the two now share a base that holds the
skeleton they agree on: a causal convolution in, a run of down blocks, a mid
block, a normalisation, a causal convolution out, and the split and gather around
them. What a family supplies is which adapter fits which down block, whether its
forward threads a temporal cache, and whether it ends on a norm that reduces over
the axis being split.

Wan's encoder adapter kept the two down block shapes it has to handle, 2.2's
grouped stages and 2.1's flat list, and gains nothing else. Unlike before it
refuses a down block it does not recognise rather than warning and leaving it
unsharded, which produced latents no one had checked.

WanEncoderAdapter had no test; both are now covered at one, two and four ranks,
including an attention among the down blocks, a non-square image, a row count
that does not divide by the rank count, and the chunked convolution path.

Co-authored-by: Cursor <cursoragent@cursor.com>

…ncoders

Both encoders are their own decoders read backwards, so the residual blocks, mid
blocks and causal convolutions are already sharded by the adapters written for the
decode side. What is new is the down block and the downsampler it holds.

Neither downsampler needs more than its convolution sharded. HunyuanVideo's
strides, and 1.5's folds each pair of rows and columns into channels, and both
read one input position per output one, so a rank can do the rest to its own rows
as long as it holds whole pairs of them. Cutting bands in whole multiples of what
the encoder narrows by is what guarantees that, and is also why the assertion
inside WanZeroPadConv2d that a band is even still holds.

HunyuanVideo ends on a GroupNorm, which reduces over the axis being split and so
is wrapped; 1.5 ends on an RMS norm, which reduces over channels and is left
alone. Covered at one, two and four ranks, with and without the attention in
HunyuanVideo's mid block, over a non-square image, a row count that does not
divide by the rank count, and the chunked convolution path.

Co-authored-by: Cursor <cursoragent@cursor.com>

The last of the causal families, and the least work: LTX-2's mid block has no
attention, so nothing has to be gathered, and its spatial padding already sits
inside the convolution rather than being applied around it.

Its down block is the one place a checkpoint has a real choice. Three of the four
downsample kinds fold space into channels after a stride-1 convolution, and the
fourth is a bare strided convolution, which only a stage that does not widen can
hold. Both routes are handled and both are tested, along with the reflection
padding LTX-2 ships with and the zeros LTX-2.3 uses.

Co-authored-by: Cursor <cursoragent@cursor.com>

Two tests per run were failing on EADDRINUSE, a different pair each time. The port fixture was
picking from 29800-39799, which overlaps the range Linux hands out to outgoing connections, so a
port a test was about to bind could be one the kernel had just given to something else.

Picking from 20000-29999 puts every port below that range. The hash is crc32 rather than hash(),
whose seed changes per interpreter, so a test now gets the same port every run and a failure can
be reproduced. Neither rules out two tests colliding with each other, and no port can be held
open for them because rank 0 has to bind the rendezvous socket itself, so run_distributed retries
on a fresh port for that.

Co-authored-by: Cursor <cursoragent@cursor.com>

AutoencoderKL's encoder, which Flux.2 builds as well, is the last of the seven this repo has a
decoder adapter for to get one for the other half. Flux-Kontext, Qwen-Image-Edit and FLUX.2 all
encode an image at the size they are about to generate at, which is the encode worth splitting.

It splits over the down blocks alone. They carry the image at full size and are the reason to
encode in parallel; the mid block and the normalisation after it run at a eighth of that, so the
split is undone before them and they run whole on every rank. That is the 2D decoder adapter read
backwards, which splits after its mid block rather than before, and it costs a mid block on every
rank in exchange for needing nothing said about the attention inside it.

Downsample2DAdapter is the part that reaches across the split. Told to pad by hand, as these
encoders tell it, the downsampler pads (0, 1, 0, 1) in its own forward and then strides over the
result, and a rank's bottom row is only padding if it is the bottom row of the whole image. The
zero-padding convolution Wan's resample already used is exactly that pair, so it moves out of the
Wan adapter to be shared, and the adapter runs the norm and the convolution itself in that case
rather than delegating, so the pad is not applied twice.

The scale factor is counted off the blocks rather than taken from the caller, because a 2D VAE
does not record its ratio anywhere and a caller reaching for the usual 8 on a three-stage encoder
would cut bands a later stage halves into a row the rank does not own.

Also corrects two descriptions of the encoder split that still described the padding the ragged
split replaced.

Co-authored-by: Cursor <cursoragent@cursor.com>

What we tune here is a property of the adapter stack rather than of the
weights: PatchGroupNorm issues the same collectives whether its input came
from Flux.2 or from torch.randn. What has to be real is the shape of the
work, and that lives in a VAE's config.json, so the architecture can be
built with random weights in a second and measured on real GPUs without
downloading anything.

The counts are the point. An optimisation that removes an all_reduce shows
up as an integer, not as a timing delta the size of the run-to-run noise on
a consumer GPU, so a change like collapsing PatchGroupNorm's three
collectives into one can land as an assertion rather than as a benchmark.
Latency, peak memory and the single-rank equivalence check come along with
it, the last being the invariant every change in here has to preserve.

Attribution is by call site, read with sys._getframe rather than by walking
the stack, which is cheap enough to leave enabled during a timed decode.

This cannot speak for real activation distributions: random weights give a
mean near zero and a variance near one, which is the easy case for any
variance computation, so a change whose error depends on the mean being
large relative to the spread still needs a real decode to sign off.

Three ways this measured something other than what a runner model does.

Importing xfuser puts AITER's GroupNorm in torch.nn.GroupNorm's place, and
both the adapter and xDiT's selection ask isinstance against whichever class
is bound when they ask. Building the VAE before that import left it holding
the class from before the swap, so a decoder made entirely of the blocks the
2D adapter wants was reported as fitting no adapter at all. xDiT then puts
the stock class back when it validates --use_parallel_vae, because AITER's
carries no num_channels for GroupNormAdapter to read; a VAE built here rather
than by a runner model has to be walked through both steps by hand.

Rank 0 alone computed the reference while the others went on to a barrier.
That barrier was the first collective in the process, so it was also where
the communicator got built, and the other three sat in init - not in the
barrier - until the store gave up ten minutes later. The reference is now
taken on every rank, which the matching seeds make free, and only at sizes
where one card can hold the whole half; above that it is skipped with a note,
since an unsharded decode at that size is the thing sharding exists to avoid.

Selection now has to come from xDiT. Choosing an adapter here when xfuser was
missing measured this file's opinion of which one fits, and reported the
mismatch as an assertion from inside a half-replaced decoder.

Agreement is reported against the reference's own scale. On random weights an
absolute tolerance says little, and in bf16 one step at magnitude 1 is
already 0.008.

PatchGroupNorm made three round trips per call and the first of them moved
a single number: the rows this rank holds, all-reduced to the group's total
so nelements can be worked out. That sum and the group sums are reductions
over the same ranks and neither depends on the other, so they can share one
tensor. The arithmetic is untouched - the same two totals come back, and a
row count is exact in float32 - which is worth saying because the third
reduction is not like this: it takes the squares about the mean the second
one produces, and folding it in would mean changing the estimator.

Measured at a 64x64 latent on four GPUs, this takes a Flux.2 decode from 75
all-reduces to 50, and the total from 163 collectives to 138. The payload
was never the point: all 75 moved 0.01 MB between them, so what this buys is
25 fewer launches and 25 fewer points where the ranks have to meet.

exchange_halo offered its bottom halo to the next rank, then blocked in the
receive from the previous one, and only once that had landed did it offer its
top halo back the other way. The two directions do not depend on each other,
so the second exchange was waiting on the first for no reason: two exposed
round trips per sharded convolution where one would do.

All four operations are now built as P2POps and issued together. NCCL groups
a batch into one operation, which also settles a complaint it was making
about the old shape - "An unbatched P2P op (send/recv) was called on this
ProcessGroup with size 4. In lazy initialization mode, this will result in a
new 2-rank NCCL communicator to be created" - once per op, and a Flux.2
decode at a 64x64 latent issues fifty-six of them.

The receive-buffer caching is unchanged, just lifted into a helper now that
both sides want it.

The bench harness has to wrap distributed_c10d as well as the re-export from
torch.distributed: P2POp validates the op it is handed against that module's
own isend and irecv, so counting collectives would otherwise have turned this
exchange into an invalid op.

…e dtype sets

use_uniform_patch built the patch boundaries from a rank's own patch size
instead of gathering everyone's, on the assumption that all patches are the
same size. Patchify cuts bands in whole multiples of the VAE's scale factor and
gives the ranks that come first one extra band where the count does not divide,
so that assumption is false for any row count that is not a multiple of the
rank count, and the boundaries it derives are wrong for every rank past the
remainder. Its cache made that worse by keying on the local patch size alone,
which does not determine the layout it stands for.

Both adapter entry points already passed False, so nothing reachable used it;
what was left was a parameter threaded through thirteen files and one test that
exercised the wrong path. The all_gather it was meant to avoid is worth
avoiding, but it needs the row layout carried forward from Patchify rather than
guessed at, which is a separate change.

PatchAdaGroupNorm goes with it: nothing has referenced it, and it would raise
on any modern torch, calling torch.tensor on a list of tensors and asking
sum_to_size to reduce (N, C, H, W) to (N, groups).

Three gaps, all of which hid work we have never measured.

--half encoder sharded the encoder and then timed vae.decode, so the encoder
adapters have never been through RCCL here at all. The half under test now
decides what it is handed and what is called, and the reference threshold is
read in latent space either way so one number means the same thing for both.

The table held the two 2D families. It now holds the five the adapters claim to
support, with the shapes taken from each checkpoint's own vae/config.json
rather than shrunk: Wan, Qwen-Image, HunyuanVideo, HunyuanVideo 1.5 and LTX-2.
Four of those carry a frame axis, so --frames sizes it, and the compression
ratios and latent width are stated in the table because they are readable off a
built VAE under three different names depending on the class.

Counts came from rank 0, which borders one neighbour where the middle ranks
border two, and so sends and receives less of a halo than they do. Every
collective is one they all wait on, so the report now carries the most any rank
made alongside rank 0's, and the per-rank spread.

smoke_families.py builds all seven on the meta device, which catches a config
key the installed diffusers does not take before it costs a pod.

At stride 1 every term that mentions where a patch sits cancels out of the halo
width, so the all_gather each convolution made to learn the other ranks patch
sizes was buying an answer the kernel already gave. Skipping it removes 143 of
the 288 collectives a Wan encode makes and 143 of 178 per-conv gathers on the
decode. A strided conv still gathers: its halo turns on where in the global
stride grid its patch begins, which is not local knowledge.

The metadata tuple now carries global_start rather than the whole boundary list,
because that is all three callers ever read from it, and it is None exactly when
no one paid to find it out.

calc_bottom_halo_width asserts its way out of a patch narrower than the kernel
reaches, so the [3,2,2]-with-a-7-kernel corner of the parametrisation had no
gathered answer to agree with. Patchify refuses that split long before a
convolution sees it.

…me port twice

The parametrisation [2,-2] failed the gate as 'process 1 terminated with signal
SIGABRT', with 'terminate called without an active exception' and both ranks
reporting a clean Gloo connect. That is a rank tearing its context down while
its peer still holds one, not a wrong answer, and the same case passed on an
earlier run of the same commit. A barrier before the teardown makes the ranks
leave in step.

The port was also derived from hash(nodeid), which python salts per
interpreter, so a failing test bound a different port on every run and could
not be asked to fail again. crc32 keeps the per-test uniqueness the fixture was
written for and makes it reproducible.

…order

A comparison of parallel VAE against tiling against a narrowed window had no home: the harness
sharded unconditionally and never tiled, so three of the four arms could only be had from a full
model run. Tiling turns out to need nothing from a runner - vae_tiling reads a diffusers VAE and
nothing else - so the runner sequence transplants whole, including the part that matters, which
is reading the VAE's own tile area before --vae_tile_size narrows it.

--no-parallel-vae gives the unsharded baseline the other arms are read against, and the docstring
now says what the harness's peak VRAM is and is not: the VAE's own, never a run's.

…harded decode

The four-arm table wants crossing with resolution, and one pod per cell makes that ninety-odd
pods for what is a few minutes of actual measurement: the wall clock is startup, the branch
install, importing torch, and building the VAE, none of which the second cell needs to pay
again. --grid-arms and --grid-shapes cross the two in one process instead.

The VAE is still rebuilt per cell, because sharding and the batched decode both replace parts of
it in place and unpicking that reliably is harder than paying for a fresh one from a fixed seed.
What is reused is the reference, which depends only on the shape and is the expensive part.

A cell that throws no longer takes the grid down with it. The ranks vote on each cell before any
of them moves on, because a rank carrying into the next cell's collectives while the others
unwind an exception would hang the pod rather than lose a row.

…fit from

A world_size 1 VAE group still ran every collective on the decode path. Measured on
flux2's decoder at 1 GPU: 52 of them per decode, about 5% of wall clock, all answering
questions the rank could answer alone.

  PatchGroupNorm   two all_reduce per norm, ~50 per decode. Summing one rank's numbers
                   across one rank returns them unchanged, so both are the identity.
  gather_patches   two all_gather, one for the sizes and one for the payload. The only
                   patch is the local one and the only size is its own.
  _init_rank_mapping  one all_gather_object per adapter constructed, since initialize()
                   clears the cache each time. A one-rank group's mapping is [this rank].

Arithmetic is untouched at every rank count: the guards skip reductions that are already
no-ops rather than computing anything differently, so the ws>1 paths are byte-identical
and ws==1 now returns what the collectives would have returned.

Note this does NOT make one rank use the chunked convolution path - PatchConv2d/3d still
short-circuit to a plain F.conv at world_size 1, ignoring block_size. Letting one rank
chunk without a halo is a separate memory win and a larger change.

Co-authored-by: Cursor <cursoragent@cursor.com>

Wrap Diffusers upsampling modules and decoder blocks in place so stale copied implementations and their unused imports can be removed.

Co-authored-by: Cursor <cursoragent@cursor.com>

Keep the suite focused on observable behavior by dropping source archaeology, one-time migration assertions, and superseded manual scripts.

Co-authored-by: Cursor <cursoragent@cursor.com>

Use the tile height for row sharding so narrow rectangular windows are not rejected by their width.

Co-authored-by: Cursor <cursoragent@cursor.com>

Keep both latent axes available so row sharding, quality bounds, and benchmark area accounting each use the dimension they actually mean.

Co-authored-by: Cursor <cursoragent@cursor.com>

Make tile marking and latent-row inspection part of the supported public boundary consumed by xDiT.

Co-authored-by: Cursor <cursoragent@cursor.com>

Count rows and columns with their own strides so the ordering test also covers non-square grids correctly.

Co-authored-by: Cursor <cursoragent@cursor.com>

Expose row count and processing factor so integrations can add operation-specific remediation without parsing error text.

Co-authored-by: Cursor <cursoragent@cursor.com>

Clarify how PatchGroupNorm aggregates shard statistics and that its forward path is inference-only.

Co-authored-by: Cursor <cursoragent@cursor.com>

Document why Patchify clones narrow views to preserve rank-local storage and contiguity.

Co-authored-by: Cursor <cursoragent@cursor.com>

Cap deterministic move and swap evaluation so first-time high-resolution layouts do not spend seconds hill-climbing negligible load differences.

Co-authored-by: Cursor <cursoragent@cursor.com>